package com.fg.commons.cryptography;

import org.apache.commons.collections.iterators.IteratorEnumeration;
import org.apache.commons.io.IOUtils;

import javax.servlet.Filter;
import javax.servlet.FilterChain;
import javax.servlet.FilterConfig;
import javax.servlet.ServletException;
import javax.servlet.ServletRequest;
import javax.servlet.ServletResponse;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpSession;
import java.io.IOException;
import java.io.InputStream;
import java.io.PrintWriter;
import java.security.KeyPair;
import java.util.Collections;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.Map;

/**
 * This filter envelopes logic connected with jCryption request parameters encryption. Basically it does:
 *
 * 1) returns jcryption.js file on urls:
 *    - /jCryption/jquery.jcryption.js
 *    - /jCryption/jquery.jcryption.min.js
 * 2) generates private + public key on url:
 *    - /jCryption/generateKeyPair
 *    and public key sends to the browser, private key stores into the session
 * 3) looks for jCryption parameter in HTTP request and when parameters is found tries to decrypt passed data with
 *    last private key stored in session (private key is for single use only and then is removed)
 *
 * Filters and servlets further in the chain works with HttpServletRequestWrapper, that transparently provides both
 * original parameters and decrypted ones as if they all comes in original request.
 *
 * @author Michal Franc, Jan Novotný, FG Forrest, donated to the www.jcryption.org
 * @version $Id: JCryptionFilter.java,v 1.5 2011/05/18 09:31:21 u_novoj Exp $
 */
public class JCryptionFilter implements Filter {
	private static final int KEY_SIZE = 512;
	private static final String SESSION_KEY = "__jCryption__keyPair";
	private static final String JCRYPTION_PARAMETER_NAME = "jCryption";
	private static final String JCRYPTION_DECRYPTED = "__jCryption_decrypted";

	private int keySize = KEY_SIZE;
	private String cryptedParameterName = JCRYPTION_PARAMETER_NAME;
	private final jCryption jCryptionInstance = new jCryption();

	@Override
	public void init(FilterConfig filterConfig) throws ServletException {
		String keySizeParam = filterConfig.getInitParameter("keySize");
		if (keySizeParam != null) {
			keySize = Integer.parseInt(keySizeParam);
		}
		if (filterConfig.getInitParameter("cryptedParameterName") != null) {
			cryptedParameterName = filterConfig.getInitParameter("cryptedParameterName");
		}
	}

	@Override
	public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain) throws IOException, ServletException {
		if (request instanceof HttpServletRequest && response instanceof HttpServletResponse) {
			HttpServletRequest httpRequest = (HttpServletRequest)request;
			HttpServletResponse httpResponse = (HttpServletResponse)response;
			//noinspection IfStatementWithTooManyBranches
			if (httpRequest.getRequestURI().endsWith("/jCryption/jquery.jcryption.js")) {
				//supply jcryption.js from classpath
				copyStreamToOutput(httpResponse, "jquery.jcryption-1.1.js");
				//request is processed, don't continue with chain processing
				return;
			} else if (httpRequest.getRequestURI().endsWith("/jCryption/jquery.jcryption.min.js")) {
				//supply minified jcryption.js from classpath
				copyStreamToOutput(httpResponse, "jquery.jcryption-1.1.min.js");
				//request is processed, don't continue with chain processing
				return;
			} else if (httpRequest.getRequestURI().endsWith("/jCryption/generateKeyPair")) {
				//generate keys and spit them to the output
				generateKeys(httpRequest, httpResponse, jCryptionInstance, keySize);
				//request is processed, don't continue with chain processing
				return;
			} else if (httpRequest.getParameterMap().containsKey(cryptedParameterName)) {
				//when crypted parameter is noticed
				HttpSession session = httpRequest.getSession();
				KeyPair keys = (KeyPair)session.getAttribute(SESSION_KEY);
			    String encrypted = request.getParameter(cryptedParameterName);
			    //decrypt data with last provided key stored in session
			    if ( encrypted!=null && keys!=null )   {
					String data = jCryptionInstance.decryptWithRedundancyCheck(encrypted, keys);
					//remove key from session - it's for single use only
					session.removeAttribute(SESSION_KEY);
					if (data != null) {
						Map<String, Object> params = jCryption.parse(data, httpRequest.getCharacterEncoding());
						//store data in request attributes so that multiple calls of the filter take advantage of
						//already decrypted data
						httpRequest.setAttribute(JCRYPTION_DECRYPTED, params);
						//wrap the request and pretend that crypted data came as ordinary POST data
						request = new JCryptionHttpServletRequest(httpRequest, params, cryptedParameterName);
					}
			    }
			} else if (httpRequest.getAttribute(JCRYPTION_DECRYPTED) != null) {
				//use already decrypted and parsed data from request
				@SuppressWarnings({"unchecked"})
				Map<String, Object> params = (Map<String, Object>)httpRequest.getAttribute(JCRYPTION_DECRYPTED);
				request = new JCryptionHttpServletRequest(httpRequest, params, cryptedParameterName);
			}
		}
		//process filter chain on
		chain.doFilter(request, response);
	}

	@Override
	public void destroy() {
		//do nothing
	}

	/**
	 * Generates public and private keys. Public key gets stored into the session - public is sent over the wire
	 * to the browser.
	 *
	 * @param httpRequest
	 * @param httpResponse
	 * @param jCryptionInstance
	 * @param keySize
	 * @throws IOException
	 */
	private static void generateKeys(HttpServletRequest httpRequest, HttpServletResponse httpResponse, jCryption jCryptionInstance, int keySize) throws IOException {
		KeyPair keys = jCryptionInstance.generateKeypair(keySize);
		httpRequest.getSession().setAttribute(SESSION_KEY, keys);
		String exponent = jCryption.getPublicKeyExponent(keys);
		String modulus = jCryption.getPublicKeyModulus(keys);
		String maxDigits = String.valueOf(jCryption.getMaxDigits(keySize));

		PrintWriter out = httpResponse.getWriter();
		out.print("{\"e\":\"");
		out.print(exponent);
		out.print("\",\"n\":\"");
		out.print(modulus);
		out.print("\",\"maxdigits\":\"");
		out.print(maxDigits);
		out.print("\"}");
	}

	/**
	 * Streams file from the classpath to the output.
	 * Used to easily embed jCryption javascript files to the client page.
	 *
	 * @param httpResponse
	 * @param fileName
	 * @throws IOException
	 */
	private static void copyStreamToOutput(HttpServletResponse httpResponse, String fileName) throws IOException {
		InputStream jsStream = Thread.currentThread().getContextClassLoader().getResourceAsStream("META-INF/jCryption/" + fileName);
		try {
			httpResponse.setContentType("text/javascript; charset=utf-8");
			httpResponse.setHeader("Content-Disposition", "attachment; filename=" + fileName);
			httpResponse.setHeader("Content-Disposition-Type", "attachment");
			IOUtils.copy(jsStream, httpResponse.getOutputStream());
		} finally {
			IOUtils.closeQuietly(jsStream);
		}
	}

	/**
	 * This class represents HTTP servlet request wrapper, that combines original parameters and those encrypted in
	 * jCryption parameter.
	 */
	private static class JCryptionHttpServletRequest extends HttpServletRequestWrapper {
		private final Map<String, Object> combinedParameters;

		private JCryptionHttpServletRequest(HttpServletRequest request, Map<String, Object> decryptedParameters, String cryptedParameterName) {
			super(request);
			//noinspection unchecked
			this.combinedParameters = new HashMap<String, Object>(request.getParameterMap());
			this.combinedParameters.putAll(decryptedParameters);
			this.combinedParameters.remove(cryptedParameterName);
		}

		@Override
		public String getParameter(String name) {
			Object value = combinedParameters.get(name);
			return value instanceof String[] ? ((String[])value)[0] : (String)value;
		}

		@Override
		public Map getParameterMap() {
			return Collections.unmodifiableMap(combinedParameters);
		}

		@Override
		public Enumeration getParameterNames() {
			return new IteratorEnumeration(combinedParameters.keySet().iterator());
		}

		@Override
		public String[] getParameterValues(String name) {
			if (combinedParameters.containsKey(name)) {
				Object value = combinedParameters.get(name);
				return value instanceof String[] ? (String[])value : new String[] {(String)value};
			} else {
				return null;
			}
		}

	}

}
